from langgraph.graph import StateGraph, MessagesState, END
from server.agent_chat.llms import Agent


class AgentChat:

    def __init__(self):
        self.agent = Agent()

    def create_workflow(self):
        workflow = StateGraph(MessagesState)
        # Define the node
        workflow.add_node("agent", self.agent)

        # Build graph
        workflow.set_entry_point("agent")
        workflow.add_edge("agent", END)

        # Compile
        return workflow.compile()

